package hamaster.jpbc.crypt.rc4;

import hamaster.jpbc.crypt.Crypt;

public class RC4 implements Crypt {

	private byte[] key;

	@Override
	public byte[] encrypt(byte[] plain) {
		return rc4(plain, key);
	}

	@Override
	public byte[] decrypt(byte[] cipher) {
		return rc4(cipher, key);
	}

	@Override
	public void setKey(byte[] key) {
		if (key == null)
			throw new IllegalArgumentException("Empty key");
		switch (key.length) {
		case 16:
			break;
		case 24:
			break;
		case 32:
			break;
		default:
			throw new IllegalArgumentException("Incorrect key length");
		}
		this.key = key;
	}

	public final byte[] rc4(byte[] buf, byte[] key) {
		if (key == null || key.length == 0)
			throw new IllegalArgumentException("Empty key");
		if (buf == null)
			return null;
		byte state[] = new byte[256];
		int x = 0;
		int y = 0;
		int index1 = 0;
		int index2 = 0;
		int xorIndex;
		byte tmp;
		for (int i = 0; i < 256; i++)
			state[i] = (byte) i;
		byte init0;
		for (int i = 0; i < 256; i++) {
			index2 = ((key[index1] & 0xff) + (state[i] & 0xff) + index2) & 0xff;
			init0 = state[i];
			state[i] = state[index2];
			state[index2] = init0;
			index1 = (index1 + 1) % key.length;
		}
		byte[] result = new byte[buf.length];
		for (int i = 0; i < buf.length; i++) {
			x = (x + 1) & 0xff;
			y = ((state[x] & 0xff) + y) & 0xff;
			tmp = state[x];
			state[x] = state[y];
			state[y] = tmp;
			xorIndex = ((state[x] & 0xff) + (state[y] & 0xff)) & 0xff;
			result[i] = (byte) (buf[i] ^ state[xorIndex]);
		}
		return result;
	}

	@Override
	public void setBlockSize(int blockSize) {
		// 不必实现此方法
		throw new RuntimeException("No need to implement this method!");
	}

}
